import paddle
import numpy as np
from sklearn.metrics import auc, roc_curve, precision_recall_curve


class ShangHaiTechMetrics():
    def __init__(self, gt_path='list/gt-sh.npy'):
        self.gt = list(np.load(gt_path))
        self.pred = paddle.zeros([0])

    def collect(self, predict, **kwargs):
        predict = paddle.squeeze(predict["scores"], 1)
        predict = paddle.mean(predict, 0)
        self.pred = paddle.concat((self.pred, predict))

    def compute(self):
        pred = list(self.pred.numpy())
        pred = np.repeat(np.array(pred), 16)
        fpr, tpr, threshold = roc_curve(self.gt, pred)
        rec_auc = auc(fpr, tpr)
        precision, recall, th = precision_recall_curve(self.gt, pred)
        pr_auc = auc(recall, precision)
        return rec_auc, pr_auc

    def clear(self):
        self.pred = paddle.zeros([0])